import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error
import pickle



def generate_model():
    df = pd.read_csv('./data/SJZ_Property.csv')
    # 定义训练集测试集
    y = df['property_price']
    data = df.drop(labels=['Unnamed: 0','property_price'], axis=1)
    train_x, test_x, train_y, test_y = train_test_split(data, y, random_state=1)
    rf = RandomForestRegressor(n_estimators=100, n_jobs=-1, oob_score=True, bootstrap=True, random_state=42)
    rf.fit(train_x, train_y)
    pickle.dump(rf, open('./model/rf_model.pickle', 'wb'))
    predict_result = rf.predict(test_x)
    # mean_absolute_error(test_y, predict)

def predict(data):
    rf = pickle.load(open('./model/rf_model.pickle', 'rb'))
    return rf.predict(data)

